import requests
import time
import json
import sys
import os
import argparse
from datetime import datetime, timedelta
import requests

# 确保中文显示正常
import matplotlib.pyplot as plt
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]

# 配置API基础URL
BASE_URL = "http://127.0.0.1:5000/api"

# 支持的漏洞类型
SUPPORTED_VULN_TYPES = {
    "sql_injection": "SQL注入",
    "xss": "跨站脚本(XSS)",
    "info_leak": "信息泄露",
    "command_injection": "命令注入",
    "csrf": "跨站请求伪造(CSRF)",
    "security_header": "安全响应头缺失"
}

# 检查系统状态
def check_system_status():
    """检查系统各组件状态 - 在演示模式下直接返回成功"""
    # 检测是否为演示模式
    if '--demo' in sys.argv:
        print("\n⚠ 系统状态检查：由于是演示模式，跳过API连接检查")
        print("  - 整体状态: 模拟正常")
        print("  - 数据库: 模拟正常")
        print("  - Redis: 模拟正常")
        print("  - Celery: 模拟正常")
        return True
        
    try:
        # 检查健康状态API
        health_response = requests.get("http://127.0.0.1:5000/health", timeout=5)
        health_data = health_response.json()
        
        print("系统状态检查：")
        print(f"  - 整体状态: {'正常' if health_data.get('status') == 'healthy' else '异常'}")
        print(f"  - 数据库: {'正常' if health_data.get('database') == 'healthy' else '异常'}")
        print(f"  - Redis: {'正常' if health_data.get('redis') == 'healthy' else '异常'}")
        print(f"  - Celery: {'正常' if health_data.get('celery_available') == 'True' else '异常'}")
        
        return health_data.get('status') == 'healthy'
    except Exception as e:
        print(f"系统状态检查失败: {str(e)}")
        return False

# 创建扫描任务
def create_scan(target_url, scan_depth=2, scan_types=None, scheduled=False):
    """创建扫描任务"""
    try:
        # 如果未指定扫描类型，默认扫描所有支持的类型
        if scan_types is None:
            scan_types = list(SUPPORTED_VULN_TYPES.keys())
        
        # 验证目标URL格式
        if not target_url.startswith(('http://', 'https://')):
            target_url = f"http://{target_url}"
        
        # 准备扫描请求数据
        scan_data = {
            "target": target_url,
            "scan_depth": scan_depth,
            "scan_types": scan_types,
            "scheduled": scheduled
        }
        
        # 发送创建扫描任务的请求
        response = requests.post(f"{BASE_URL}/scan", json=scan_data, timeout=30)
        response.raise_for_status()
        
        result = response.json()
        task_id = result.get("task_id")
        scan_id = result.get("scan_id")
        scan_type_names = [SUPPORTED_VULN_TYPES.get(t, t) for t in scan_types]
        
        print(f"\n✓ 成功创建扫描任务!")
        print(f"  任务ID: {task_id}")
        print(f"  扫描ID: {scan_id}")
        print(f"  目标URL: {target_url}")
        print(f"  扫描深度: {scan_depth}")
        print(f"  扫描类型: {', '.join(scan_type_names)}")
        
        return task_id
    except requests.exceptions.RequestException as e:
        print(f"✗ 创建扫描任务失败: {str(e)}")
        if hasattr(e, 'response') and e.response:
            print(f"  错误详情: {e.response.text}")
        return None

# 检查扫描进度
def check_scan_progress(task_id):
    """检查扫描任务的进度 - 在演示模式下模拟进度"""
    if '--demo' in sys.argv or task_id.startswith('demo_task_'):
        # 演示模式，模拟扫描进度
        start_time = time.time()
        
        # 模拟扫描过程，耗时约10秒
        for progress in range(0, 101, 5):
            elapsed_time = int(time.time() - start_time)
            
            # 显示进度条
            bar_length = 30
            filled_length = int(bar_length * progress // 100)
            bar = '█' * filled_length + '-' * (bar_length - filled_length)
            
            # 模拟不同阶段的状态
            if progress < 20:
                status_text = '等待中'
            elif progress < 95:
                status_text = '运行中'
            else:
                status_text = '已完成'
            
            sys.stdout.write(f"\r扫描进度: [{bar}] {progress}% | 状态: {status_text} | 用时: {elapsed_time}秒")
            sys.stdout.flush()
            
            if progress == 100:
                sys.stdout.write("\n")
                return True
            
            # 根据进度调整等待时间，模拟真实扫描节奏
            if progress < 30:
                time.sleep(0.2)  # 初始快速推进
            elif progress < 70:
                time.sleep(0.4)  # 中间稳定推进
            else:
                time.sleep(0.6)  # 最后阶段较慢推进
        
        return True
        
    try:
        start_time = time.time()
        
        # 每3秒检查一次进度，直到扫描完成或超时
        max_wait_time = 300  # 最大等待时间，单位：秒
        
        while time.time() - start_time < max_wait_time:
            # 获取扫描状态
            response = requests.get(f"{BASE_URL}/scan/{task_id}/status", timeout=5)
            response.raise_for_status()
            
            status_data = response.json()
            progress = status_data.get("progress", 0)
            status = status_data.get("status", "unknown")
            
            # 计算已用时间
            elapsed_time = int(time.time() - start_time)
            
            # 显示进度条
            bar_length = 30
            filled_length = int(bar_length * progress // 100)
            bar = '█' * filled_length + '-' * (bar_length - filled_length)
            status_text = {
                'pending': '等待中',
                'running': '运行中',
                'completed': '已完成',
                'failed': '已失败'
            }.get(status, status)
            
            sys.stdout.write(f"\r扫描进度: [{bar}] {progress}% | 状态: {status_text} | 用时: {elapsed_time}秒")
            sys.stdout.flush()
            
            if status in ["completed", "failed"]:
                sys.stdout.write("\n")
                return True
            
            # 每3秒检查一次
            time.sleep(3)
        
        print(f"\n✗ 扫描超时 (超过{max_wait_time}秒)")
        return False
    except requests.exceptions.RequestException as e:
        print(f"\n✗ 获取扫描进度失败: {str(e)}")
        return False

# 获取扫描报告
def get_scan_report(task_id, detailed=True):
    """获取扫描任务的报告 - 在演示模式下生成模拟报告"""
    if '--demo' in sys.argv or task_id.startswith('demo_task_'):
        # 演示模式，生成模拟报告数据
        now = datetime.now()
        created_at = now.strftime("%Y-%m-%d %H:%M:%S")
        finished_at = (now + timedelta(seconds=15)).strftime("%Y-%m-%d %H:%M:%S")
        
        # 模拟漏洞数据
        mock_vulnerabilities = [
            {
                'vuln_id': 'VULN-2025-001',
                'vuln_type': 'sql_injection',
                'path': '/api/products?category=electronics&sort=price',
                'severity': 'high',
                'risk_score': 9.2,
                'evidence': "参数未经过滤: 'category' 参数包含可能的SQL注入模式 '1 OR 1=1'",
                'suggestion': "对所有用户输入参数进行参数化查询处理，避免直接拼接SQL语句。使用ORM框架如SQLAlchemy或Django ORM可以有效防止SQL注入攻击。"
            },
            {
                'vuln_id': 'VULN-2025-002',
                'vuln_type': 'xss',
                'path': '/search?q=test',
                'severity': 'medium',
                'risk_score': 6.5,
                'evidence': "用户输入未经过滤直接输出到页面: 'q' 参数值未进行HTML转义",
                'suggestion': "对所有用户可控输入进行适当的HTML转义处理，或使用内容安全策略(CSP)来限制JavaScript的执行范围。"
            },
            {
                'vuln_id': 'VULN-2025-003',
                'vuln_type': 'info_leak',
                'path': '/.env',
                'severity': 'high',
                'risk_score': 8.7,
                'evidence': "敏感配置文件可被未授权访问: 服务器返回了.env文件内容，包含数据库凭证",
                'suggestion': "确保所有敏感配置文件不在Web根目录下，并通过Web服务器配置限制对这些文件的访问。"
            },
            {
                'vuln_id': 'VULN-2025-004',
                'vuln_type': 'security_header',
                'path': '/login',
                'severity': 'low',
                'risk_score': 3.2,
                'evidence': "缺少安全响应头: Content-Security-Policy未设置",
                'suggestion': "为所有HTTP响应添加适当的安全头，包括Content-Security-Policy、X-Content-Type-Options、X-Frame-Options等。"
            },
            {
                'vuln_id': 'VULN-2025-005',
                'vuln_type': 'csrf',
                'path': '/user/change_password',
                'severity': 'medium',
                'risk_score': 6.8,
                'evidence': "未实现CSRF保护机制: 敏感操作请求中未包含CSRF令牌验证",
                'suggestion': "为所有涉及状态变更的表单和API请求实现CSRF令牌验证机制。"
            }
        ]
        
        # 计算风险统计
        risk_stats = {
            'total': len(mock_vulnerabilities),
            'high': sum(1 for v in mock_vulnerabilities if v['severity'] == 'high'),
            'medium': sum(1 for v in mock_vulnerabilities if v['severity'] == 'medium'),
            'low': sum(1 for v in mock_vulnerabilities if v['severity'] == 'low')
        }
        
        # 模拟报告对象
        report = {
            'scan_id': task_id.replace('task', 'scan'),
            'target': 'http://example.com',
            'status': 'completed',
            'created_at': created_at,
            'finished_at': finished_at,
            'risk_stats': risk_stats,
            'vulnerabilities': mock_vulnerabilities
        }
        
        # 格式化时间
        
        print(f"\n扫描报告 - ID: {report.get('scan_id')}")
        print(f"目标: {report.get('target')}")
        print(f"状态: {'成功' if report.get('status') == 'completed' else report.get('status')}")
        print(f"开始时间: {created_at}")
        print(f"完成时间: {finished_at}")
        print("  注意: 此报告基于模拟数据生成，用于演示目的")
        
        # 显示风险统计
        print("\n风险统计:")
        print(f"  - 总漏洞数: {risk_stats.get('total', 0)}")
        print(f"  - 高危漏洞: {risk_stats.get('high', 0)}")
        print(f"  - 中危漏洞: {risk_stats.get('medium', 0)}")
        print(f"  - 低危漏洞: {risk_stats.get('low', 0)}")
        
        # 按漏洞类型统计
        vuln_type_count = {}
        vulnerabilities = report.get('vulnerabilities', [])
        
        for vuln in vulnerabilities:
            vuln_type = vuln.get('vuln_type')
            vuln_type_count[vuln_type] = vuln_type_count.get(vuln_type, 0) + 1
        
        if vuln_type_count:
            print("\n漏洞类型分布:")
            for vuln_type, count in vuln_type_count.items():
                vuln_type_name = SUPPORTED_VULN_TYPES.get(vuln_type, vuln_type)
                print(f"  - {vuln_type_name}: {count}个")
        
        # 如果需要详细信息，显示漏洞详情
        if detailed and vulnerabilities:
            print("\n漏洞详情:")
            # 按严重程度和风险评分排序
            sorted_vulnerabilities = sorted(
                vulnerabilities,
                key=lambda x: (['high', 'medium', 'low'].index(x.get('severity', 'low')), 
                             -x.get('risk_score', 0))
            )
            
            for i, vuln in enumerate(sorted_vulnerabilities, 1):
                print(f"\n{i}. 漏洞ID: {vuln.get('vuln_id')}")
                print(f"   类型: {SUPPORTED_VULN_TYPES.get(vuln.get('vuln_type'), vuln.get('vuln_type'))}")
                print(f"   路径: {vuln.get('path')}")
                print(f"   严重程度: {vuln.get('severity')}")
                print(f"   风险评分: {vuln.get('risk_score')}")
                print(f"   证据: {vuln.get('evidence')}")
                
                # 如果有修复建议，显示简短建议
                if vuln.get('suggestion'):
                    suggestion_preview = vuln.get('suggestion')[:100] + ('...' if len(vuln.get('suggestion')) > 100 else '')
                    print(f"   修复建议: {suggestion_preview}")
            
            # 分析攻击路径
            print("\n正在分析攻击路径...")
            attack_paths = analyze_attack_paths(vulnerabilities)
            report['attack_paths'] = attack_paths
        elif not vulnerabilities:
            print("\n未发现漏洞")
        
        # 保存完整报告到文件
        report_file = f"scan_report_{task_id}.json"
        with open(report_file, 'w', encoding='utf-8') as f:
            json.dump(report, f, ensure_ascii=False, indent=2)
        print(f"\n✓ 完整报告已保存到: {report_file}")
        
        return report
        
    try:
        # 正常模式，通过API获取报告
        response = requests.get(f"{BASE_URL}/scan/{task_id}/report", timeout=10)
        response.raise_for_status()
        
        report = response.json()
        
        # 格式化时间
        created_at = format_time(report.get('created_at'))
        finished_at = format_time(report.get('finished_at'))
        
        print(f"\n扫描报告 - ID: {report.get('scan_id')}")
        print(f"目标: {report.get('target')}")
        print(f"状态: {'成功' if report.get('status') == 'completed' else report.get('status')}")
        print(f"开始时间: {created_at}")
        print(f"完成时间: {finished_at}")
        
        # 显示风险统计
        risk_stats = report.get('risk_stats', {})
        print("\n风险统计:")
        print(f"  - 总漏洞数: {risk_stats.get('total', 0)}")
        print(f"  - 高危漏洞: {risk_stats.get('high', 0)}")
        print(f"  - 中危漏洞: {risk_stats.get('medium', 0)}")
        print(f"  - 低危漏洞: {risk_stats.get('low', 0)}")
        
        # 按漏洞类型统计
        vuln_type_count = {}
        vulnerabilities = report.get('vulnerabilities', [])
        
        for vuln in vulnerabilities:
            vuln_type = vuln.get('vuln_type')
            vuln_type_count[vuln_type] = vuln_type_count.get(vuln_type, 0) + 1
        
        if vuln_type_count:
            print("\n漏洞类型分布:")
            for vuln_type, count in vuln_type_count.items():
                vuln_type_name = SUPPORTED_VULN_TYPES.get(vuln_type, vuln_type)
                print(f"  - {vuln_type_name}: {count}个")
        
        # 如果需要详细信息，显示漏洞详情
        if detailed and vulnerabilities:
            print("\n漏洞详情:")
            # 按严重程度和风险评分排序
            sorted_vulnerabilities = sorted(
                vulnerabilities,
                key=lambda x: (['high', 'medium', 'low'].index(x.get('severity', 'low')), 
                             -x.get('risk_score', 0))
            )
            
            for i, vuln in enumerate(sorted_vulnerabilities, 1):
                print(f"\n{i}. 漏洞ID: {vuln.get('vuln_id')}")
                print(f"   类型: {SUPPORTED_VULN_TYPES.get(vuln.get('vuln_type'), vuln.get('vuln_type'))}")
                print(f"   路径: {vuln.get('path')}")
                print(f"   严重程度: {vuln.get('severity')}")
                print(f"   风险评分: {vuln.get('risk_score')}")
                print(f"   证据: {vuln.get('evidence')}")
                
                # 如果有修复建议，显示简短建议
                if vuln.get('suggestion'):
                    suggestion_preview = vuln.get('suggestion')[:100] + ('...' if len(vuln.get('suggestion')) > 100 else '')
                    print(f"   修复建议: {suggestion_preview}")
        elif not vulnerabilities:
            print("\n未发现漏洞")
        
        # 保存完整报告到文件
        report_file = f"scan_report_{task_id}.json"
        with open(report_file, 'w', encoding='utf-8') as f:
            json.dump(report, f, ensure_ascii=False, indent=2)
        print(f"\n✓ 完整报告已保存到: {report_file}")
        
        return report
    except requests.exceptions.RequestException as e:
        print(f"\n✗ 获取扫描报告失败: {str(e)}")
        if hasattr(e, 'response') and e.response:
            print(f"  错误详情: {e.response.text}")
        return None

# 创建模拟扫描
def create_mock_scan():
    """创建模拟扫描任务 - 在演示模式下直接返回模拟任务ID"""
    print("\n====== 创建模拟扫描任务 ======")
    
    # 模拟数据，不依赖API
    task_id = "demo_task_" + datetime.now().strftime("%Y%m%d%H%M%S")
    scan_id = "demo_scan_" + datetime.now().strftime("%Y%m%d%H%M%S")
    
    print(f"✓ 模拟扫描任务创建成功!")
    print(f"  任务ID: {task_id}")
    print(f"  扫描ID: {scan_id}")
    print(f"  目标URL: http://example.com")
    print("  注意: 这是一个模拟扫描，会生成丰富的演示数据")
    print("  当前系统处于离线演示模式，使用内置的模拟数据")
    
    return task_id

# 分析攻击路径
def analyze_attack_paths(vulnerabilities):
    """分析漏洞之间的关联性，生成可能的攻击路径"""
    if not vulnerabilities:
        print("  没有发现漏洞，无法分析攻击路径")
        return
    
    # 按风险评分和严重程度对漏洞进行排序
    sorted_vulns = sorted(
        vulnerabilities,
        key=lambda x: (['high', 'medium', 'low'].index(x.get('severity', 'low')),
                      -x.get('risk_score', 0))
    )
    
    # 分析可能的攻击链
    attack_chains = []
    
    # 查找可能的入口点（通常是可公开访问的漏洞）
    entry_points = [v for v in sorted_vulns if v.get('path') and not v.get('path').startswith('/admin')]
    
    if entry_points:
        print("\n可能的攻击路径分析:")
        
        # 对每个入口点分析可能的攻击链
        for i, entry in enumerate(entry_points[:3], 1):  # 只显示前3个最可能的攻击路径
            chain = [entry]
            
            # 查找可能的后续漏洞（通常是信息泄露后的进一步利用）
            entry_path = entry.get('path', '')
            
            # 查找相关的后续漏洞
            for vuln in sorted_vulns:
                if vuln == entry:  # 跳过自身
                    continue
                
                # 简单的关联规则：
                # 1. 如果入口是信息泄露，可能导致更严重的漏洞利用
                # 2. 如果路径相关，可能在同一功能模块中
                vuln_path = vuln.get('path', '')
                
                # 检查关联条件
                if ((entry.get('vuln_type') == 'info_leak' and vuln.get('severity') == 'high') or
                    (entry_path.split('/')[1] == vuln_path.split('/')[1] if len(entry_path.split('/')) > 1 and len(vuln_path.split('/')) > 1 else False)):
                    chain.append(vuln)
                    break
            
            # 显示攻击链
            print(f"\n攻击路径 {i}:")
            
            for step, vuln in enumerate(chain, 1):
                vuln_type = SUPPORTED_VULN_TYPES.get(vuln.get('vuln_type'), vuln.get('vuln_type'))
                print(f"  步骤 {step}: 通过 {vuln_type} 漏洞 ({vuln.get('path')})")
                print(f"     - 严重程度: {vuln.get('severity')}")
                print(f"     - 风险评分: {vuln.get('risk_score')}")
                
                if step == len(chain):
                    # 显示最终可能的影响
                    if vuln.get('severity') == 'high':
                        impact = "可能导致系统完全被控制、数据泄露或服务中断"
                    elif vuln.get('severity') == 'medium':
                        impact = "可能导致数据泄露或服务异常"
                    else:
                        impact = "可能导致有限的信息泄露"
                    print(f"     - 潜在影响: {impact}")
                else:
                    print(f"     - 可被利用来: 访问更深层功能或获取敏感信息")
        
        # 提供修复建议的优先级
        print("\n修复优先级建议:")
        for i, vuln in enumerate([v for v in sorted_vulns if v.get('severity') == 'high'][:3], 1):
            print(f"  {i}. 立即修复: {SUPPORTED_VULN_TYPES.get(vuln.get('vuln_type'), vuln.get('vuln_type'))} - {vuln.get('path')}")
    else:
        print("  未能识别明确的攻击入口点")

# 格式化时间
def format_time(time_str):
    """格式化时间字符串为更易读的形式"""
    try:
        # 处理不同格式的时间字符串
        if 'Z' in time_str:
            # ISO格式带Z
            dt = datetime.fromisoformat(time_str.replace('Z', '+00:00'))
        else:
            dt = datetime.fromisoformat(time_str)
        
        # 转换为本地时间并格式化
        local_time = dt.astimezone()
        return local_time.strftime("%Y-%m-%d %H:%M:%S")
    except:
        # 如果解析失败，返回原始字符串
        return time_str

# 扫描单个URL
def scan_single_url(url, scan_depth=2, scan_types=None):
    """扫描单个URL"""
    print(f"\n开始扫描URL: {url}")
    
    # 创建扫描任务
    task_id = create_scan(url, scan_depth=scan_depth, scan_types=scan_types)
    if not task_id:
        return False
    
    print("\n✓ 扫描任务已成功创建!")
    print("正在等待扫描完成...")
    
    # 检查扫描进度
    if check_scan_progress(task_id):
        print("\n扫描已完成，正在获取扫描报告...")
        
        # 获取并显示扫描报告
        report = get_scan_report(task_id, detailed=True)
        
        # 注意：在get_scan_report中已包含攻击路径分析
        if report and 'attack_paths' in report:
            print("\n攻击路径分析已完成")
        elif report and 'vulnerabilities' in report and report['vulnerabilities']:
            print("\n正在分析攻击路径...")
            analyze_attack_paths(report['vulnerabilities'])
        
        return True
    else:
        print("\n✗ 扫描失败或超时")
        return False

# 交互式扫描
def interactive_scan():
    """交互式扫描模式"""
    print("\n进入交互式扫描模式")
    
    # 获取目标URL
    target_url = input("请输入要扫描的目标URL (例如: https://example.com): ").strip()
    if not target_url:
        print("URL不能为空")
        return False
    
    # 获取扫描深度
    try:
        scan_depth = int(input("请输入扫描深度 (1-3, 默认2): ").strip() or "2")
        scan_depth = max(1, min(3, scan_depth))  # 限制在1-3之间
    except ValueError:
        scan_depth = 2
        print("使用默认扫描深度: 2")
    
    # 选择扫描类型
    print("\n请选择要扫描的漏洞类型 (输入编号，用逗号分隔，默认全部):")
    print("0: 全部漏洞类型")
    
    vuln_type_list = list(SUPPORTED_VULN_TYPES.items())
    for i, (vuln_type, vuln_name) in enumerate(vuln_type_list, 1):
        print(f"{i}: {vuln_name}")
    
    selection = input("请选择: ").strip()
    if not selection or selection == '0':
        scan_types = None  # 全部类型
    else:
        try:
            # 解析用户选择
            selected_indices = [int(idx.strip()) for idx in selection.split(',')]
            scan_types = []
            
            for idx in selected_indices:
                if 1 <= idx <= len(vuln_type_list):
                    scan_types.append(vuln_type_list[idx-1][0])
                else:
                    print(f"警告: 无效的选择 {idx}，已忽略")
            
            if not scan_types:
                scan_types = None  # 如果没有有效选择，使用全部类型
        except ValueError:
            scan_types = None
            print("输入无效，使用全部漏洞类型")
    
    # 开始扫描
    return scan_single_url(target_url, scan_depth=scan_depth, scan_types=scan_types)

# 批量扫描
def batch_scan(file_path):
    """批量扫描文件中的URL"""
    try:
        # 读取URL列表
        with open(file_path, 'r', encoding='utf-8') as f:
            urls = [line.strip() for line in f if line.strip()]
        
        if not urls:
            print(f"文件 {file_path} 中没有有效的URL")
            return False
        
        print(f"\n开始批量扫描，共 {len(urls)} 个URL")
        
        # 扫描每个URL
        for i, url in enumerate(urls, 1):
            print(f"\n{'-'*50}")
            print(f"扫描 URL {i}/{len(urls)}: {url}")
            print(f"{'-'*50}")
            
            # 使用默认参数扫描
            scan_single_url(url, scan_depth=1)
        
        return True
    except Exception as e:
        print(f"批量扫描失败: {str(e)}")
        return False

# 主函数
def main():
    # 设置命令行参数解析器
    parser = argparse.ArgumentParser(description="漏洞扫描系统命令行工具")
    parser.add_argument('--url', help='要扫描的目标URL')
    parser.add_argument('--file', help='包含URL列表的文件路径，用于批量扫描')
    parser.add_argument('--depth', type=int, default=2, help='扫描深度 (1-3)')
    parser.add_argument('--interactive', action='store_true', help='启用交互式扫描模式')
    parser.add_argument('--demo', action='store_true', help='运行演示模式（使用模拟数据）')
    
    args = parser.parse_args()
    
    print("=" * 60)
    print("          漏洞扫描系统演示工具          ")
    print("=" * 60)
    print("注意: 请仅在授权的系统上使用此工具进行安全测试")
    print("      未经授权对系统进行扫描可能违反法律法规")
    print("=" * 60)
    
    # 1. 检查系统状态
    print("\n1. 检查系统状态...")
    if not check_system_status():
        print("系统状态异常，无法继续扫描")
        proceed = input("是否继续？(y/n): ")
        if proceed.lower() != 'y':
            return
    
    # 2. 根据命令行参数决定扫描模式
    success = False
    
    if args.interactive:
        # 交互式扫描模式
        success = interactive_scan()
    elif args.file:
        # 批量扫描模式
        success = batch_scan(args.file)
    elif args.url:
        # 单URL扫描模式
        success = scan_single_url(args.url, scan_depth=args.depth)
    elif args.demo:
        # 演示模式
        print("\n运行演示模式...")
        task_id = create_mock_scan()
        if task_id:
            # 监控扫描进度
            if check_scan_progress(task_id):
                # 获取并显示扫描报告
                get_scan_report(task_id)
                success = True
    else:
        # 默认模式：显示帮助信息并提供示例
        print("\n使用方法:")
        print("  python demo_scan.py --interactive      # 启动交互式扫描")
        print("  python demo_scan.py --url https://example.com  # 扫描指定URL")
        print("  python demo_scan.py --file urls.txt    # 批量扫描文件中的URL")
        print("  python demo_scan.py --demo            # 运行演示模式")
        print("\n示例:")
        # 执行一个简单的示例扫描
        print("\n执行演示扫描...")
        task_id = create_mock_scan()
        if task_id:
            # 监控扫描进度
            if check_scan_progress(task_id):
                # 获取并显示扫描报告
                get_scan_report(task_id)
                success = True
    
    print("\n" + "=" * 60)
    print(f"扫描任务{'完成' if success else '失败'}")
    print("=" * 60)
    print("\n提示:")
    print("1. 您也可以通过Web界面(http://localhost:3001)进行操作")
    print("2. 对于真实环境的扫描，请确保您有目标系统的授权")
    print("3. 扫描结果会同时存储在数据库和Redis缓存中，提高查询效率")
    print("4. 扫描报告已保存为JSON文件，可以用文本编辑器查看完整内容")

if __name__ == "__main__":
    main()